# -*- coding: utf-8 -*- #

# -----------------------------------------------------------------------
# File Name:    test.py
# Version:      ver1_0
# Created:      2024/06/17
# Description:  本文件定义了模型的测试流程
#               ★★★请在空白处填写适当的语句，将模型测试流程补充完整★★★
# -----------------------------------------------------------------------

import torch
from torch.utils.data import DataLoader
from torchvision.transforms import ToTensor
from dataset import CustomDataset
from model import CustomNet as MyModel


def test(dataloader, model, device):
    model.eval()
    size = len(dataloader.dataset)
    correct_num = 0
    test_loss = 0.1

    # 添加调试信息
    sample_batch = next(iter(dataloader))
    print(f"输入图像尺寸: {sample_batch['image'].shape}")

    with torch.no_grad():
        for batch in dataloader:
            images = batch['image'].to(device)
            labels = batch['label'].to(device)

            # 添加形状检查
            print(f"批量图像尺寸: {images.shape}")

            outputs = model(images)
            _, predicted = torch.max(outputs.data, 1)
            correct_num += (predicted == labels).sum().item()

    accuracy = 100 * correct_num / size
    print(f'测试准确率: {accuracy:.2f}%')
    return accuracy

if __name__ == "__main__":
    # 加载训练好的模型
    model = MyModel()
    #model = torch.load('./models/best_model.pth')
    model.load_state_dict(torch.load('./models/final_model.pth'))

    if torch.cuda.is_available():
        device = torch.device("cuda")
    else:
        device = torch.device("cpu")
    model.to(device)


    # 测试数据加载器
    test_dataloader = DataLoader(CustomDataset('./images/test.txt', './images/test', ToTensor),
                                 batch_size=32)
    # 运行测试函数
    test(test_dataloader, model, device)
